import os

from langchain.agents import AgentExecutor
from langchain_community.chat_models import QianfanChatEndpoint
from qianfan.extensions.langchain.agents import QianfanSingleActionAgent

from THUDM.tool.db_employee_tool import query_employee

os.environ["QIANFAN_AK"] = "ZGgmCMM95MfabqYLG2swwRWM"
os.environ["QIANFAN_SK"] = "yDAIOpvEu7vm2phjYDT5sdIg5h4LBb4v"

llm = QianfanChatEndpoint(model="ERNIE-Bot")

tools = [query_employee]

agent = QianfanSingleActionAgent.from_system_prompt(tools, llm)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

while True:
    # 查询用户admin的信息
    input_str = input("请输入您的问题：")
    result = agent_executor(input_str)
    print("执行结果: ")
    print(result)
